from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import pandas as pd

# 加载数据集
wine = load_wine()
X = wine.data
y = wine.target
feature_names = wine.feature_names
target_names = wine.target_names

# 划分训练集和测试集 (8:2)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

# 标准化数据（KNN对量纲敏感，朴素贝叶斯的高斯分布假设也需要标准化）
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# 查看预处理后的数据形状
print(f"训练集样本数: {X_train_scaled.shape[0]}, 特征数: {X_train_scaled.shape[1]}")
print(f"测试集样本数: {X_test_scaled.shape[0]}")